import sys

import tensorflow as tf
import SimpleITK as sitk

from Model import *
from Utils import *

config = dict()

# config for inputs
config['input_shape']         = (96, 96, 80)
config['batch_size']          = 1

# config for network
config['network']             = 'mynetwork'
config['weight_decay']        = False

# config for training
config['iterations']          = 20000
config['learning_rate']       = 1e-5
config['learning_rate_decay'] = False

def connectedComponentFilter(array, isFirst=False, threshold=512):
    connectedFilter = sitk.ConnectedComponentImageFilter()
    connectedFilter.SetDebug(False)
    connectedFilter.SetFullyConnected(False)

    tempImage = sitk.GetImageFromArray(array.astype(np.uint8))
    filteredImage = connectedFilter.Execute(*[tempImage])
    tempResult = sitk.GetArrayFromImage(filteredImage)

    btm = False
    for i in range(1, tempResult.max() + 1):
        s = (tempResult == i).sum()
        if s > threshold:
            if np.argwhere(tempResult == i)[:, 2].min() == 0:
                btm = True

    if not isFirst:
        btm = False

    max = -1
    maxLabel = -1
    for i in range(1, tempResult.max() + 1):
        s = (tempResult == i).sum()
        if s < threshold or (btm == True and np.argwhere(tempResult == i)[:, 2].min() != 0):
            continue
        elif s > max:
            max = s
            maxLabel = i

    ret =  (tempResult == maxLabel).astype(np.uint8)
    return ret

def centerOfSegment(array, ROICenter):
    indexs = np.argwhere(array)
    center = indexs.sum(0) / indexs.shape[0]
    return [int(center[i]) + ROICenter[i] for i in range(3)]

def run(image, model_path):
    dims = image.shape
    bboxSize = [96, 96, 80]
    delta = [bboxSize[i] // 2 for i in range(3)]
    center = [0, 0, 0]

    # normalization
    nRange = [-500.0, 2500.0]
    image = (image - nRange[0]) / (nRange[1] - nRange[0])
    image = np.clip(image, 0, 1)

    # extend image
    extendSize = [dims[i] + bboxSize[i] + 2 for i in range(3)]
    imageExtend = np.zeros(extendSize)
    memoryExtend = np.zeros(extendSize)
    imageExtend[:dims[0], :dims[1], :dims[2]] = image

    # tensor
    inputPatch = np.zeros(np.hstack([1, bboxSize, 2]), dtype=np.float32)

    net = Unet(config)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver = tf.train.Saver()
        saver.restore(sess, model_path)

        # Segment first vertebra
        print("Start finding first vertebrae.")
        strides = [(dims[i] - bboxSize[i] - 1) /
                   delta[i] + 2 for i in range(3)]
        axes = np.mgrid[0:strides[2], 0:strides[1], 0:strides[0]]
        z, y, x = [axis.flatten().astype(np.int16) for axis in axes]
        for stride in zip(x, y, z):
            center = [stride[i] * delta[i] for i in range(3)]
            print(center)
            inputPatch[0, ..., 0] = slide(imageExtend, center, bboxSize)
            prediction = net.predict(sess, inputPatch)[0, ..., 0]
            prediction = connectedComponentFilter(prediction, True)
            if prediction.sum() > 0:
                center = centerOfSegment(prediction, center)
                break

        vSize, pSize = [], []
        for i in range(1, 25):
            print("Segment", i)
            for innerIter in range(3):
                center = [max(0, center[idx] - delta[idx]) for idx in range(3)]
                inputPatch[0, ..., 0] = slide(imageExtend, center, bboxSize)
                temp = slide(memoryExtend, center, bboxSize).copy()
                temp[temp > 0] = 1
                inputPatch[0, ..., 1] = temp
                prediction = net.predict(sess, inputPatch)[0, ..., 0]
                prediction = connectedComponentFilter(prediction)
                prediction = prediction - inputPatch[0, ..., 1] * prediction

                print(center, dims[2])
                lastCenter = center[:]
                center = centerOfSegment(prediction, center)

            if prediction.sum() == 0:
                break

            upper = np.argwhere(prediction)[:, 2].max() + center[2] - delta[2]
            memoryExtend[lastCenter[0]: lastCenter[0] + bboxSize[0], lastCenter[1]: lastCenter[1] +
                         bboxSize[1], lastCenter[2]: lastCenter[2] + bboxSize[2]] += i * prediction
            vSize.append(prediction.sum())
            pSize.append(prediction.max(0).sum())

            if upper >= dims[2] - 1:
                break
            if i > 3 and vSize[-2] - vSize[-3] > 1000 and pSize[-2] - pSize[-3] > 100 and vSize[-3] < 15000:
                break

        return memoryExtend[:dims[0], :dims[1], :dims[2]]


if __name__ == '__main__':
    fileName = sys.argv[1]
    outputName = sys.argv[2]
    model_path = sys.argv[3]

    reader = sitk.ImageFileReader()
    reader.SetFileName(fileName)
    inputImage = reader.Execute()

    inputSize = inputImage.GetSize()
    inputSpacing = inputImage.GetSpacing()
    inputOrigin = inputImage.GetOrigin()
    inputDirection = inputImage.GetDirection()
    resolution = [1.25, 1.25, 1.25]
    newSize = (np.array(inputSize) * np.array(inputSpacing) /
               np.array(resolution) + 0.5).astype(np.int)

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(inputImage)
    resampler.SetOutputSpacing(resolution)
    resampler.SetSize(newSize.tolist())

    inputImage = resampler.Execute(inputImage)
    array = sitk.GetArrayFromImage(inputImage)
    array = np.swapaxes(array, 0, 2)
    result = run(array, model_path)
    result = np.swapaxes(result, 0, 2)

    tempImage = sitk.GetImageFromArray(result)
    tempImage.SetDirection(inputDirection)
    tempImage.SetOrigin(inputOrigin)
    tempImage.SetSpacing(resolution)

    resampler = sitk.ResampleImageFilter()
    resampler.SetInterpolator(sitk.sitkNearestNeighbor)
    resampler.SetReferenceImage(tempImage)
    resampler.SetOutputSpacing(inputSpacing)
    resampler.SetSize(inputSize)
    outputImage = resampler.Execute(tempImage)

    castFilter = sitk.CastImageFilter()
    castFilter.SetOutputPixelType(sitk.sitkInt8)
    outputImage = castFilter.Execute(label_smoothing(outputImage))

    write_image(outputImage, outputName)
